/*
 * Copyright (C) 2021-2022 Red Hat, Inc.
 *
 * Written By: Vadim Rozenfeld <vrozenfe@redhat.com>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met :
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and / or other materials provided with the distribution.
 * 3. Neither the names of the copyright holders nor the names of their contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include "pch.h"
#include "strsafe.h"

constexpr DWORD RESTART_DELAY_MS = 3000;
constexpr DWORD SC_ACTION_COUNT = 3;

extern CService srvc;

void ShowUsage()
{
    LOGI("USAGE:");
    LOGI("vmgtsvr -i: Install service");
    LOGI("vmgtsvr -u: Uninstall service");
    LOGI("vmgtsvr -r: Run service");
    LOGI("vmgtsvr -s: Stop service");
    LOGI("vmgtsvr status: Current status");
}

BOOL InstallService()
{
    SC_HANDLE newService;
    SC_HANDLE scm;
    TCHAR szBuffer[255];
    TCHAR szPath[MAX_PATH];

    GetModuleFileName( GetModuleHandle(NULL), szPath, MAX_PATH );
    LOGI(L"Nodule Name: {}", GET_PRINTABLE_STRING(szPath));
    if (FAILED( StringCchCopy(szBuffer, 255, TEXT("\"")))) {
        LOGE(L"szBuffer: {}", GET_PRINTABLE_STRING(szBuffer));
        return FALSE;
    }
    if (FAILED( StringCchCat(szBuffer, 255, szPath))) {
        LOGE(L"szBuffer: {}", GET_PRINTABLE_STRING(szBuffer));
        return FALSE;
    }
    if (FAILED( StringCchCat(szBuffer, 255, TEXT("\"")))) {
        LOGE(L"szBuffer: {}", GET_PRINTABLE_STRING(szBuffer));
        return FALSE;
    }

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_CREATE_SERVICE);
    if (scm == NULL) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
        return FALSE;
    }
    newService = CreateService(scm, ServiceName, DisplayName, SERVICE_ALL_ACCESS,
                               SERVICE_WIN32_OWN_PROCESS | SERVICE_INTERACTIVE_PROCESS,
                               SERVICE_AUTO_START, SERVICE_ERROR_NORMAL, szBuffer,
                               NULL, NULL, NULL, NULL, NULL);
    if (!newService) {
        LOGE("Failed to call CreateService with error: {}", GetLastError());
        CloseServiceHandle(scm);
        return FALSE;
    } else {
        LOGI("Service Installed");
        ServiceRun();
        if (!SetServiceFailureActions(newService)) {
            LOGE("Failed to set service failure actions");
        }
    }

    CloseServiceHandle(newService);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL SetServiceFailureActions(SC_HANDLE service)
{
    if (!service) {
        LOGE("service is a null pointer");
        return FALSE;
    }

    SC_ACTION actions[SC_ACTION_COUNT] = {
        // Set up the first restart action
        {SC_ACTION_RESTART, RESTART_DELAY_MS},
        // Set up the second restart action
        {SC_ACTION_RESTART, RESTART_DELAY_MS},
        // Set up the subsequent restart action
        {SC_ACTION_RESTART, RESTART_DELAY_MS}
    };

    SERVICE_FAILURE_ACTIONS serviceFailureActions;
    serviceFailureActions.dwResetPeriod = INFINITE;
    serviceFailureActions.lpRebootMsg = nullptr;
    serviceFailureActions.lpCommand = nullptr;
    serviceFailureActions.cActions = SC_ACTION_COUNT;
    serviceFailureActions.lpsaActions = actions;

    if (!ChangeServiceConfig2(service, SERVICE_CONFIG_FAILURE_ACTIONS, &serviceFailureActions)) {
        LOGE("Failed to call ChangeServiceConfig2 with error: {}", GetLastError());
        return FALSE;
    }

    return TRUE;
}

BOOL UninstallService()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    SERVICE_STATUS status;

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS | DELETE);
    if (!service) {
        LOGE("Failed to call OpenService with error: {}", GetLastError());
    }

    res = QueryServiceStatus(service, &status);
    if (!res) {
        LOGE("Failed to call QueryServiceStatus with error: {}", GetLastError());
    }

    if (status.dwCurrentState != SERVICE_STOPPED) {
        LOGI("Stopping service...");
        res = ControlService(service, SERVICE_CONTROL_STOP, &status);
        if (!res) {
            LOGE("Failed to call ControlService with error: {}", GetLastError());
        }
        Sleep(5000);
    }

    res = DeleteService(service);
    if (res) {
        LOGI("Service Uninstalled");
    } else {
        LOGE("Failed to call DeleteService with error: {}", GetLastError());
    }

    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL ServiceRun()
{
    SC_HANDLE scm, Service;
    SERVICE_STATUS ssStatus;
    DWORD dwOldCheckPoint;
    DWORD dwStartTickCount;
    DWORD dwWaitTime;
    DWORD dwStatus;

    LOGI("Run Service");

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
    }

    Service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!Service) {
        LOGE("Failed to call OpenService with error: {}", GetLastError());
        return FALSE;
    } else {
        StartService(Service, 0, NULL);
        srvc.GetStatus(Service);

        if (!QueryServiceStatus( Service, &ssStatus) ) {
            LOGE("Failed to call QueryServiceStatus with error: {}", GetLastError());
        }
        dwStartTickCount = GetTickCount();
        dwOldCheckPoint = ssStatus.dwCheckPoint;

        while (ssStatus.dwCurrentState == SERVICE_START_PENDING) {
            dwWaitTime = ssStatus.dwWaitHint / 10;

            if( dwWaitTime < 1000 ) {
                dwWaitTime = 1000;
            } else if ( dwWaitTime > 10000 ) {
                dwWaitTime = 10000;
            }

            Sleep( dwWaitTime );

            if (!QueryServiceStatus(Service, &ssStatus) ) {
                LOGE("Failed to call QueryServiceStatus with error: {}", GetLastError());
                break;
            }

            if ( ssStatus.dwCheckPoint > dwOldCheckPoint ) {
                dwStartTickCount = GetTickCount();
                dwOldCheckPoint = ssStatus.dwCheckPoint;
            } else {
                if(GetTickCount()-dwStartTickCount > ssStatus.dwWaitHint) {
                    break;
                }
            }
        }

        if (ssStatus.dwCurrentState == SERVICE_RUNNING) {
            srvc.GetStatus(Service);
            dwStatus = NO_ERROR;
        } else {
            LOGI("Service not started.");
            LOGI("Current State: {}", ssStatus.dwCurrentState);
            LOGI("Exit Code: {}", ssStatus.dwWin32ExitCode);
            LOGI("Service Specific Exit Code: {}", ssStatus.dwServiceSpecificExitCode);
            LOGI("Check Point: {}", ssStatus.dwCheckPoint);
            LOGI("Wait Hint: {}", ssStatus.dwWaitHint);
            dwStatus = GetLastError();
            LOGE("ServiceRun failed with error: {}", dwStatus);
        }
    }

    CloseServiceHandle(scm);
    CloseServiceHandle(Service);

    return TRUE;
}

BOOL ServiceControl(int ctrl)
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res = TRUE;
    SERVICE_STATUS status;

    LOGI("Service control with value: {}", ctrl);

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!service) {
        LOGE("Failed to call OpenService with error: {}", GetLastError());
    }

    if (ctrl == SERVICE_CONTROL_STOP) {
        LOGI("Service is stopping...");
        res = ControlService(service, SERVICE_CONTROL_STOP, &status);
    }

    if (!res) {
        LOGE("Failed to call ControlService with error: {}", GetLastError());
    } else {
        srvc.GetStatus(service);
    }

    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}

BOOL GetConfiguration()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    LPQUERY_SERVICE_CONFIG buffer;
    DWORD sizeNeeded;

    LOGI("Get configuration");

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (!scm) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
    }

    service = OpenService(scm, ServiceName, SERVICE_QUERY_CONFIG);
    if (!service) {
        LOGE("Failed to call OpenService with error: {}", GetLastError());
    }

    buffer = (LPQUERY_SERVICE_CONFIG)LocalAlloc(LPTR, 4096);
    if (!buffer) {
        LOGE("Failed to call LocalAlloc with error: {}", GetLastError());
    }

    res = QueryServiceConfig(service, buffer, 4096, &sizeNeeded);
    if (!res) {
        LOGE("Failed to call QueryServiceConfig with error: {}", GetLastError());
    }

    LOGI(L"Service name: {}", buffer->lpDisplayName);
    LOGI("Service type: {}", buffer->dwServiceType);
    LOGI("Start type: {}", buffer->dwStartType);
    LOGI(L"Start name: {}", buffer->lpServiceStartName);
    LOGI(L"Path: {}", GET_PRINTABLE_STRING(buffer->lpBinaryPathName));

    LocalFree(buffer);

    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}


BOOL ChangeConfig()
{
    SC_HANDLE service;
    SC_HANDLE scm;
    BOOL res;
    SC_LOCK lock;

    LOGI("Change Config");

    scm = OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS | GENERIC_WRITE);
    if (!scm) {
        LOGE("Failed to call OpenSCManager with error: {}", GetLastError());
    }
    lock = LockServiceDatabase(scm);
    if (lock == 0) {
        LOGE("Failed to call LockServiceDatabase with error: {}", GetLastError());
    }
    service = OpenService(scm, ServiceName, SERVICE_ALL_ACCESS);
    if (!service) {
        LOGE("Failed to call OpenService with error: {}", GetLastError());
    }
    res = ChangeServiceConfig(
                             service,
                             SERVICE_NO_CHANGE,
                             SERVICE_NO_CHANGE,
                             SERVICE_NO_CHANGE,
                             NULL,
                             NULL,
                             NULL,
                             NULL,
                             NULL,
                             NULL,
                             NULL
                             );
    if (!res) {
        UnlockServiceDatabase(lock);
        LOGE("Failed to call ChangeServiceConfig with error: {}", GetLastError());
    }

    res = UnlockServiceDatabase(lock);
    if (!res) {
        LOGE("Failed to call UnlockServiceDatabase with error: {}", GetLastError());
    }
    CloseServiceHandle(service);
    CloseServiceHandle(scm);

    return TRUE;
}
